import xml.etree.ElementTree as ET
import os
import glob

# 定义类别，这里定义自己的类别，有多少类按照顺序写上
# classes = ["C_CLSDXXJ", 'C_JCXDXXJ', '0', '1', '2', '3', '4']
classes = ['C_DG', 'C_BHT_DT']


# 将标注框转换为YOLO格式
def convert(size, box):
    dw = 1.0 / size[0]
    dh = 1.0 / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


# 转换 XML 标注文件为 YOLO 格式
def convert_annotation(image_name):
    # 构造 XML 和 TXT 文件路径
    xml_path = os.path.join('/media/liuhw/38185FB1185F6CBE/disk_d/DSG/4_1/xml_out', image_name[:-3] + 'xml')
    txt_path = os.path.join('/media/liuhw/38185FB1185F6CBE/disk_d/DSG/4_1/txt_out', image_name[:-3] + 'txt')

    # 如果标签输出文件夹不存在，创建它
    os.makedirs(os.path.dirname(txt_path), exist_ok=True)

    # 打开并读取 XML 文件
    try:
        with open(xml_path, 'r') as f:
            xml_text = f.read()
            root = ET.fromstring(xml_text)
    except FileNotFoundError:
        print(f"XML 文件未找到: {xml_path}")
        return

    # 获取图片的宽度和高度
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    # 打开输出的 TXT 文件进行写入
    with open(txt_path, 'w') as out_file:
        # 遍历每个 object 标签
        for obj in root.iter('object'):
            cls = obj.find('name').text
            # 如果该类别不在指定类别列表中，跳过
            if cls not in classes:
                print(f"跳过类别: {cls}")
                continue
            cls_id = classes.index(cls)

            # 获取标注框信息
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
                 float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
            bb = convert((w, h), b)

            # 将转换后的结果写入输出文件
            out_file.write(f"{cls_id} {' '.join([str(a) for a in bb])}\n")


# 主程序
if __name__ == '__main__':
    # 遍历所有的 .jpg 图片文件，并转换成相应的 .xml 文件
    for image_path in glob.glob("/media/liuhw/38185FB1185F6CBE/disk_d/DSG/4_1/img_out/*.jpg"):  # 图片路径
        image_name = os.path.basename(image_path)  # 获取图片文件名
        convert_annotation(image_name)  # 调用转换函数
